// nbody.c
// multi-thread design using OpenMP

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>
#include <sys/time.h>

#include <omp.h>

#define NUM_THREADS 4
#define CHUNK N/NUM_THREADS

typedef struct {
  float m, x, y, z, vx, vy, vz;
} p_type;

typedef struct {
  float ax, ay, az;
} a_type;

int N=81920;
float EPS=0;
int LOOP=1;

void error(char *msg);

int main(int argc, char *argv[]) {

  FILE *fp;

  p_type *p;
  a_type *a;

  int step;
  int i, j;

  int chunk;

  float rx, ry, rz;
  float d, dd;
  float s;

  struct timeval tv1, tv2, tv3, tv4;
  long long tv_comp, tv_total;

	if (argc==1) {
		printf("usage: %s <input file> [N] [EPS] [LOOP]\n", argv[0]);
		return 0;
	}

  if (argc>2) N = atoi(argv[2]);
  if (argc>3) sscanf(argv[3], "%f", &EPS);
  if (argc>4) LOOP = atoi(argv[4]);

  printf("nbody : %d particles with EPS=%f in %d iterations\n", N, EPS, LOOP);

  gettimeofday(&tv3, NULL);

  p = (p_type *)malloc(sizeof(p_type)*N);
  if (!p) error("malloc (p)");
  a = (a_type *)malloc(sizeof(a_type)*N);
  if (!a) error("malloc (a)");

  fp = fopen(argv[1], "r");
  if (fread(p, sizeof(float)*7, N, fp) < N) error("fread");
  fclose(fp);

  gettimeofday(&tv1, NULL);

  // ----------
  chunk = CHUNK;
  omp_set_num_threads(4);
  for (step=0; step<LOOP; step++) { // in each iteration
    memset(a, 0, sizeof(float)*N*3);
#pragma omp parallel default(shared) private(i, j, rx, ry, rz, d, dd, s)
    {
      if (omp_get_thread_num() == 0)
        printf("In %d parallel threads.\n", omp_get_num_threads());
#pragma omp for schedule(dynamic,chunk) nowait
    for (i=0; i<N; i++) {
      for (j=0; j<N; j++) {
        if (j!=i) { // if not self compute acceleration vector

          rx = p[j].x - p[i].x;
          ry = p[j].y - p[i].y;
          rz = p[j].z - p[i].z;

          dd = rx*rx + ry*ry + rz*rz + EPS;
          d = 1 / sqrtf(dd * dd * dd);

          s = p[j].m * d;

          a[i].ax += rx * s;
          a[i].ay += ry * s;
          a[i].az += rz * s;

        } // if j!=i
      } // for j
    } // for i

    } // omp section

    for (i=0; i<N; i++) { // for every particle
      // update position
      p[i].x += p[i].vx;
      p[i].y += p[i].vy;
      p[i].z += p[i].vz;

      // update velocity
      p[i].vx += a[i].ax;
      p[i].vy += a[i].ay;
      p[i].vz += a[i].az;
    }

  } // for step
  // ----------

  gettimeofday(&tv2, NULL);

  fp = fopen("result_nbody.dat", "w");
  if (fwrite(p, sizeof(float)*N*7, 1, fp) < 1)
    error("fwrite");
  fclose(fp);

  free(p);
  free(a);

  gettimeofday(&tv4, NULL);

  tv_comp = (tv2.tv_sec - tv1.tv_sec) * 1000000 + (tv2.tv_usec - tv1.tv_usec);
  tv_total = (tv4.tv_sec - tv3.tv_sec) * 1000000 + (tv4.tv_usec - tv3.tv_usec);
  printf("compu time = %lld usec\n", tv_comp);
  printf("total time = %lld usec\n", tv_total);

  return 0;
}

void error(char *msg) {
  printf("Err: %s\n", msg);
  exit(0);
}
